import os
import sys
import json
from datetime import datetime
from typing import Any, Dict, List, Optional, Union
import sys

# 配置导入路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import LOG_DIR, VERBOSE_LOGGING

class ModelLogger:
    """记录模型输入和输出的工具类"""
    
    def __init__(self, log_dir: str = LOG_DIR, verbose: bool = VERBOSE_LOGGING):
        self.log_dir = log_dir
        self.verbose = verbose
        self.current_session = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_path = os.path.join(log_dir, f"model_log_{self.current_session}.jsonl")
        
        # 创建日志目录
        os.makedirs(log_dir, exist_ok=True)
        
        print(f"📝 日志记录器初始化: 详细模式={'开启' if verbose else '关闭'}")
        print(f"📄 日志文件路径: {self.log_path}")
    
    def log_model_interaction(self, 
                              agent_name: str, 
                              inputs: Dict[str, Any], 
                              prompt, 
                              output: str,
                              step: str,
                              file_path: Optional[str] = None) -> None:
        """记录模型输入和输出到日志文件"""
        timestamp = datetime.now().isoformat()
        
        # 仅保留输入的前1000字符用于显示
        truncated_inputs = {}
        for k, v in inputs.items():
            if isinstance(v, str) and len(v) > 1000:
                truncated_inputs[k] = v[:1000] + "... [截断]"
            else:
                truncated_inputs[k] = v
                
        # 用于显示的简短输入
        display_inputs = {k: (v[:50] + "..." if isinstance(v, str) and len(v) > 50 else v) 
                         for k, v in inputs.items()}
        
        # 详细日志记录
        log_entry = {
            "timestamp": timestamp,
            "agent": agent_name,
            "step": step,
            "file_path": file_path,
            "inputs": truncated_inputs,  # 使用截断的输入避免日志文件过大
            "prompt": prompt,
            "output": output
        }
        
        # 将日志写入文件
        with open(self.log_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(log_entry, ensure_ascii=False) + "\n")
        
        # 如果启用了详细模式，则在控制台打印
        if self.verbose:
            border = "=" * 80
            print(f"\n{border}")
            print(f"🤖 Agent: {agent_name} | 步骤: {step} | 文件: {file_path or '无'}")
            print(f"📥 输入参数: {json.dumps(display_inputs, ensure_ascii=False, indent=2)}")
            print("📤 输出内容: ")
            print(f"{output[:200]}..." if len(output) > 200 else output)
            print(f"{border}\n")

# 全局日志记录器实例
logger = ModelLogger()
